import numpy as np
import matplotlib.pyplot as plt
import os
from scipy import io
import seaborn as sns
import pandas as pd

# equ 数据
# grades = np.array([1, 2, 3, 4, 5, 6, 7])
# training_loss = np.array([3.72937274e+00, 1.50875223e+00, 2.15167546e-07, 5.87520477e-09, 3.04074455e-09, 2.88082669e-09, 2.83508594e-09])
# validation_loss = np.array([3.25761299e+00, 1.70479580e+00, 2.76565882e-07, 8.24055998e-09, 6.45127665e-09, 6.39424694e-09, 6.83932461e-09])
# accuracy = np.array([8.86373162e-01, 6.53302491e-01, 2.70473189e-04, 8.91304080e-05, 8.67473354e-05, 8.68671123e-05, 8.95678761e-05])

# inc 数据
grades = np.array([1, 2, 3, 4, 5, 6, 7])
training_loss = np.array([2.98719406e+00, 2.23097301e+00, 1.11419886e-05, 9.44659106e-09, 6.58114807e-09, 2.69897793e-09, 2.90098612e-09])
validation_loss = np.array([3.25726652e+00, 1.98038480e+00, 1.09324827e-05, 1.21209263e-08, 7.57849372e-09, 4.63005634e-09, 1.86515985e-07])
accuracy = np.array([8.86342049e-01, 7.02935696e-01, 1.64477457e-03, 8.76287086e-05, 7.83501600e-05, 7.48204475e-05, 3.25789800e-04])

# 绘图
plt.figure(figsize=(10, 6))

plt.plot(grades, training_loss, marker='o', label='Training Loss')
plt.plot(grades, validation_loss, marker='o', label='Validation Loss')
plt.plot(grades, accuracy, marker='o', label='RE')

plt.yscale('log')  # 使用对数刻度
plt.xlabel('Grade')
plt.ylabel('Loss')
#plt.title('Training Loss, Validation Loss, and Accuracy by Grade')
plt.legend()
plt.grid(True)
# 保存图片
plt.savefig('../fig/exp1_grade2.png', dpi=300)  # 保存为高清图片

# 显示图形
plt.show()


plt.figure()
result = np.load('../mat_file/exp1_inc_256_1e-2_1.npy')
case_basename = 'exp1.mat'
rev_filepath = os.path.join('../dataset', case_basename)
mat_file = io.loadmat(rev_filepath)
real_u = mat_file['real_y'].astype(np.float32)
img_u = mat_file['img_y'].astype(np.float32)
u = np.hstack([real_u, img_u])
num_test = mat_file['num_test'][0, 0]

xx = 2 * np.pi * np.fft.fftshift(np.fft.fftfreq(num_test, 2 / (num_test-1)))
standard = np.abs(np.fft.fftshift(np.fft.fft(u[:, 0] + 1j * u[:, 1]) * 2 / num_test))

# 遍历 result 的前10行
for i in range(6):
    # 计算相对误差
    relative_error = np.abs(np.fft.fftshift(
                    np.fft.fft(u[:, 0] + 1j * u[:, 1]) * 2 / num_test - np.fft.fft(
                    result[i, :].reshape([-1]) * 2 / num_test)) / standard)

    # 绘制相对误差图
    plt.plot(xx, relative_error)

# 添加图例和坐标轴标签
plt.legend(['Grade 1', 'Grade 2', 'Grade 3', 'Grade 4', 'Grade 5',
            'Grade 6'],
           loc='lower left', bbox_to_anchor=(0.0, 0.0))
plt.xlabel('Angular Frequency(rad/s)')
plt.ylabel('Relative Error in Frequency Domain')

# 设置坐标轴范围和刻度类型
plt.xlim(-2000, 2000)
plt.yscale('log')

# 保存图像
plt.savefig('../fig/exp1_re2.png', bbox_inches='tight')

# 显示图形
plt.show()

# heatmap_data_list = []

# # 遍历 result 的前 10 行
# for i in range(10):
#     # 截取频率范围为 -600 到 600 的数据
#     result_trimmed = result[i, :]

#     # 计算傅里叶变换结果的幅度矩阵
#     fft_result = np.abs(np.fft.fftshift(np.fft.fft(result_trimmed * 2 / num_test)))
#     # 计算相对值并添加到列表中
#     heatmap_data_list.append(fft_result / standard)

# # 将列表转换为 NumPy 数组
# heatmap_data = np.array(heatmap_data_list)
# freq_range = np.logical_and(xx >= -600, xx <= 600)

# # 创建 DataFrame
# df = pd.DataFrame(data=heatmap_data[3:, freq_range])

# # 绘制热图
# plt.figure()
# sns.heatmap(df, cmap='viridis', cbar=True)

# # 设置坐标轴标签和标题
# plt.xlabel('Angular Frequency(rad/s)')
# plt.ylabel('Grade')
# plt.title('Heatmap of Fourier Transform Magnitude')

# # 保存图像
# plt.savefig('../fig/mm_seaborn.png', bbox_inches='tight')

# # 显示图形
# plt.show()






############.  final compare

# plt.figure()
# result1 = np.load('../mat_file/equal_width_net.npy')
# result2 = np.load('../mat_file/increase_width_net.npy')
# result3 = np.load('../mat_file/resnet_256_6_0.npy')
# result4 = np.load('../mat_file/multi_scale_0.npy')
#
# result = np.zeros((4, result1.shape[1]), dtype=complex)
# result[0, :] = result1[8, :]
# result[1, :] = result2[6, :]
# result[2, :] = result3[0, :]
# result[3, :] = result4[0, :]
#
# noise_level = 0
# case_basename = 'exp3_com_%.3f.mat' % (noise_level)
# rev_filepath = os.path.join('../dataset', case_basename)
# mat_file = io.loadmat(rev_filepath)
# real_u = mat_file['real_y'].astype(np.float32)
# img_u = mat_file['img_y'].astype(np.float32)
# u = np.hstack([real_u, img_u])
# num_test = mat_file['num_test'][0, 0]
#
# xx = 2 * np.pi * np.fft.fftshift(np.fft.fftfreq(num_test, 2 / (num_test-1)))
# standard = np.abs(np.fft.fftshift(np.fft.fft(u[:, 0] + 1j * u[:, 1]) * 2 / num_test))
#
# # 遍历 result 的前10行
# for i in range(4):
#     # 计算相对误差
#     relative_error = np.abs(np.fft.fftshift(
#                     np.fft.fft(u[:, 0] + 1j * u[:, 1]) * 2 / num_test - np.fft.fft(
#                     result[i, :].reshape([-1]) * 2 / num_test)) / standard)
#
#     # 绘制相对误差图
#     plt.plot(xx, relative_error)
#
# # 添加图例和坐标轴标签
# plt.legend(['E-AMGDL', 'I-AMGDL', 'RESNet', 'multi-scale'],
#            loc='lower left', bbox_to_anchor=(0.0, 0.0))
# plt.xlabel('Angular Frequency(rad/s)')
# plt.ylabel('Relative Error in Frequency Domain')
#
# # 设置坐标轴范围和刻度类型
# plt.xlim(-2000, 2000)
# plt.yscale('log')
#
# # 保存图像
# plt.savefig('../fig/compare_model.png', bbox_inches='tight')
